# -*- coding: utf-8 -*-
import pgzrun
import numpy as np
import time

WIDTH = 500
HEIGHT = 500

BACK_COLOR = (0, 0, 0)

fromColor = [0, 100, 100]
toColor = [0, 200, 200]

#生成z坐标，axis为起始位置，nx，ny为x向和y向的格点个数
def genZ(axis,n):
    x0, x1, y0, y1 = axis
    x = np.linspace(x0, x1, n)
    y = np.linspace(y0, y1, n)
    real, img = np.meshgrid(x, y)
    z = real + img * 1j
    return z

# pp = 0.45
# qq = -0.1428
# pp = 0.285
# qq = 0.01
pp = 0
qq = 0

def genC(pp, qq, n):
    pb = [pp] * n
    qb = [qq] * n
    p = np.array(pb)
    q = np.array(qb)    
    # p = np.linspace(0.25, 0.5, n)
    # q = np.linspace(-0.1428, -0.1, n)

    real, img = np.meshgrid(p, q)
    c = real + img * 1j
    return c

#获取Julia集，n为迭代次数，m为判定发散点，大于1即可
def getJulia(z, c, n, m = 2):
    t = time.time()
    c = np.zeros_like(z) + c
    out = abs(z)
    for i in range(n):
        absz = abs(z)
        z[absz>m] = 0		#对开始发散的点置零
        c[absz>m] = 0		
        out[absz>m] = i	#记录发散点的发散速度
        z = z * z + c
    print("time:",time.time() - t)
    return out
count = 50
axis = np.array([-2, 1, -1.5, 1.5])
s = 500
z0 = genZ(axis,s)
c0 = genC(pp, qq, s)
step0 = (z0 - c0) / count
cc = c0
mBrot = getJulia(z0, c0, 100)
print(mBrot.shape)
brot = mBrot.tolist()
offset_c = (s - 500) * 200 // 500
offset_r = (s - 500) * 250 // 500

ppp = 0.0
qqq = 0.0
enableGo = False

def on_key_down():
    global enableGo
    enableGo = True
    print('on_key_down')

def on_mouse_down(pos):
    print(pos)


def update():
    global enableGo, count
    if enableGo == False:
        return

    global mBrot, brot, pp, qq, z0, cc
    cc += step0    
    mBrot = getJulia(z0, cc, 100)
    brot = mBrot.tolist()
    count -= 1
    if count < 0:
        enableGo = False

def draw():
    screen.clear()
    screen.fill(BACK_COLOR)
    
    for r in range(HEIGHT):
        for c in range(WIDTH):
            clr = brot[r + offset_r][c + offset_c]
            a = [i * clr / 100 for i in toColor]
            b = [i * (1 - clr / 100) for i in fromColor]
            clr = [a[i] + b[i] for i  in range(3)]
            draw_pt((c, r), 1, tuple(clr))

def draw_pt(pos, r, clr):
    screen.draw.filled_circle(pos, r, clr)

pgzrun.go()